Skip to content

feat(memory): port memory manager and extraction to Python#2740

Merged
JackYPCOnline merged 11 commits into
strands-agents:mainfrom
JackYPCOnline:feat/memory-manager-port
Jun 15, 2026
Merged

feat(memory): port memory manager and extraction to Python#2740
JackYPCOnline merged 11 commits into
strands-agents:mainfrom
JackYPCOnline:feat/memory-manager-port

Conversation

@JackYPCOnline

@JackYPCOnline JackYPCOnline commented Jun 11, 2026

Copy link
Copy Markdown
Contributor

Description

Ports the memory module from the TypeScript SDK (strands-ts/src/memory) to the Python SDK (strands-py). The module gives agents cross-session recall and persistence through a MemoryManager plugin that manages pluggable MemoryStore backends, exposes search_memory / add_memory tools, and runs automatic background extraction that distills conversation turns into durable memory.

This is a behavior-preserving port: the TypeScript test suite is the reference, and every it(...) case has a Python counterpart. The notable design work is adapting the TS promise/event-loop model to Python asyncio.

Public API — new package strands.memory:

from strands import Agent
from strands.memory import MemoryManager

# A store implements `search` and (optionally) `add` / `add_messages` / `get_tools`,
# and may declare an extraction config to auto-distill turns into memory.
memory_manager = MemoryManager(
    stores=[my_store],
    add_tool_config=True,            # expose the add_memory tool (opt-in)
    flush_on_invocation_end=True,    # await pending extraction writes per invocation
)
agent = Agent(model=model, plugins=[memory_manager])
agent("Remember I prefer dark mode")

# Programmatic access (coroutines):
results = await memory_manager.search("user preferences")
await memory_manager.flush()

Also adds AggregateMemoryError to strands.types.exceptions — a Python 3.10-safe stand-in for JS AggregateError (ExceptionGroup is 3.11+) used to surface multi-store write failures with each underlying reason.

Asyncio adaptation (key deviation from TS). TS relies on a long-lived event loop, so fire-and-forget background saves survive between turns. Python's synchronous Agent(...) entry point runs each invocation in a fresh loop (asyncio.run), which cancels in-flight tasks on return. The port uses asyncio.gather(return_exceptions=True) for Promise.allSettled, per-store asyncio.Task chains for serialized saves, and a tracked background-task set. The opt-in flush_on_invocation_end (default False) registers an AfterInvocationEvent hook that awaits pending writes before the per-invocation loop tears down; async callers owning a persistent loop can leave it off and call flush() at a shutdown boundary. This flag has no TS equivalent and exists solely to bridge the event-loop lifecycle difference.

Scope. This PR ports strands-ts/src/memory/ only. Concrete store backends (e.g. BedrockKnowledgeBaseStore, which lives in the separate strands-ts/src/vended-memory-stores/ module) are intended as a follow-up PR; this change ships the MemoryManager machinery and the MemoryStore protocol that those backends implement.

Related Issues

N/A

Documentation PR

No new docs required; public classes are documented via docstrings.

Type of Change

New feature

Testing

How have you tested the change? Verify that the changes do not break functionality or introduce new warnings.

Ported the three TS test files (memory-manager, extraction, model-extractor) as example-based pytest suites so coverage mirrors the TS suite case-for-case; runs green under tests/strands/memory. Adjacent plugins/tools suites were run to confirm no regressions in plugin/tool discovery.

  • I ran hatch run prepare

Checklist

  • I have read the CONTRIBUTING document
  • I have added any necessary tests that prove my fix is effective or my feature works
  • I have updated the documentation accordingly
  • I have added an appropriate example to the documentation to outline the feature, or no new docs are needed
  • My changes generate no new warnings
  • Any dependent changes have been merged and published

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

Port the strands-ts memory module to strands-py: a MemoryManager plugin with search_memory/add_memory tools, pluggable MemoryStore backends, and background extraction (ExtractionCoordinator, ModelExtractor, invocation/interval triggers).

Adapts the TS promise model to asyncio: gather(return_exceptions=True) for Promise.allSettled, per-store asyncio.Task chains for serialized saves, and a tracked background-task set. Adds an opt-in flush_on_invocation_end so extraction persists under the synchronous Agent(...) entry point, whose per-invocation event loop would otherwise cancel in-flight saves.

Adds AggregateMemoryError to types/exceptions for Python 3.10-safe multi-store failure aggregation (ExceptionGroup is 3.11+).
@github-actions github-actions Bot added enhancement New feature or request area-persistence Session management or checkpointing area-async Related to asynchronous flows or multi-threading area-hooks Features or requests that might be implementable via hooks python Pull requests that update python code strands-running labels Jun 11, 2026
Comment thread strands-py/src/strands/memory/types.py
Comment thread strands-py/src/strands/memory/memory_manager.py
Comment thread strands-py/src/strands/memory/memory_manager.py
@github-actions

Copy link
Copy Markdown
Contributor

Question — shipped store backend: This PR adds the MemoryManager machinery and the MemoryStore protocol, but no concrete store implementation. The TS side ships BedrockKnowledgeBaseStore (under strands-ts/src/vended-memory-stores/), and the design doc references an InMemoryMemoryStore for prototyping. As it stands, a user can't actually use this feature without writing their own store first, and there's no integration test exercising a real end-to-end path.

Is a concrete store (and an integ test) following in a separate PR? If so, a one-line note in the description would set expectations. If this is meant to be usable on its own, shipping at least an in-memory store would make the feature self-contained and give the suite an end-to-end anchor beyond the unit-level fakes.

@github-actions

Copy link
Copy Markdown
Contributor

Process note — API review label: This PR introduces a substantial new public API surface — a new top-level strands.memory package, the MemoryManager primitive, the MemoryStore protocol that customers implement, and ~10 new exported types/dataclasses. Per team/API_BAR_RAISING.md, a new customer-facing primitive/abstraction falls under "substantial changes," which calls for an explicitly designated API reviewer (or team consensus) rather than standard PR review alone.

The PR description does an excellent job documenting use cases, signatures, and the asyncio deviation — that's exactly the proposer prep the doc asks for. Could you add the needs-api-review label so this gets the appropriate sign-off before merge? Also worth confirming the divergence from designs/0011-memory-manager.md is intentional (the design doc describes store_memory/injection/EvictionTrigger; this port mirrors the current TS surface — add_memory, no injection — which seems correct, but the design doc status is still "Proposed").

Comment thread strands-py/src/strands/memory/memory_manager.py
@github-actions

Copy link
Copy Markdown
Contributor

Assessment: Comment

A clean, well-documented port. Tests pass (106), ruff and mypy are clean, the implementation faithfully mirrors the current TS surface, and the asyncio adaptation is thoughtfully reasoned in both the code and the PR description. My feedback is concentrated on the event-loop lifecycle defaults and process, not the porting work itself.

Review themes
  • API review process: Substantial new public surface (new package + customer-facing primitive/protocol) — should carry needs-api-review per team/API_BAR_RAISING.md.
  • Asyncio lifecycle defaults: flush_on_invocation_end=False makes the most common (sync) path silently drop extraction writes; flush() also doesn't drain add-tool fire-and-forget tasks. These are the highest-impact items — the obvious path should be the happy path.
  • Public-type fidelity: MemoryManagerConfig claims to mirror the constructor but omits flush_on_invocation_end.
  • Usability/scope: No concrete MemoryStore ships, so the feature isn't self-contained; worth confirming a backend + integ test is following.
  • Minor: divergent stores=[] semantics between the tool and programmatic API.

Nicely structured, well-tested work — the whole-object equality assertions and case-for-case TS parity make the suite easy to trust.

Move the _flatten_reasons helper into memory_manager.py (where the TS source defines it) and delete memory/_errors.py, which had no TypeScript counterpart. AggregateMemoryError remains in strands.types.exceptions.
Add flush_on_invocation_end to MemoryManagerConfig so the doc/typing aid matches the constructor, and document the intentional empty-stores divergence (mirrored from TS) between the tools and the programmatic search/add APIs. Addresses PR strands-agents#2740 review.
@github-actions

Copy link
Copy Markdown
Contributor

Assessment: Comment (approve — code is solid; design follow-ups are maintainer/author calls)

Reviewed 20a5a5e6 ("address review nits in coordinator and docstrings"), which responds to @pgrayy's review. The changes are clean and the nits are well-handled:

  • process() → sync returning asyncio.Task | None (per the memory-optimization nit) — verified there are only two callers, schedule() and the test suite, and both were updated. schedule() no longer double-wraps in create_task, and the tests now route through a small _drive() helper that awaits the returned task. No stale await process(...) remains.
  • Dropped the assert extraction is not None in favor of a graceful if extraction is None: return. Good call: I confirmed the coordinator is only ever constructed from _extraction_stores, which is pre-filtered on extraction is not None (memory_manager.py:195), so the branch is genuinely unreachable in production. Tightening the field type (the alternative suggested) isn't viable since MemoryStore.extraction must stay | None for non-extraction stores — the early return is the right resolution.
  • _write now takes extractor as an explicit param rather than re-deriving it and re-asserting — removes the second assert and reads more directly.
  • _block_kind no longer special-cases "text" — correct, every content block is a single-key mapping, so the generic path already covers it.
  • Docstring trims (the snake_case note, the ExceptionGroup/3.10 rationale) — agreed these were implementation/translation context that didn't belong in shipped docstrings.

Re-verified on 20a5a5e6: 108 memory tests pass, ruff and mypy clean.

Still open — but these are maintainer/author decisions, not code blockers

The substantive design questions from the review are tracked in-thread and don't block the port:

  • Structured output for the extractor prompt (reliability)
  • Possible under-extraction when a fact spans turns trimmed before extraction
  • Whether to add a dedicated memory_manager param on Agent.__init__ (symmetry with conversation_manager/session_manager)
  • Auto-enabling flush_on_invocation_end on sync invoke
  • MemorySearchToolConfig naming symmetry
  • Recording the "not true fire-and-forget" caveat on the background-tool-task ticket
  • Integ tests — author noted these land after Owen's PR

From my side the code is in good shape and all actionable review nits on this commit are addressed; the remaining items are design conversations for the team.

@opieter-aws opieter-aws enabled auto-merge (squash) June 12, 2026 21:14
@opieter-aws opieter-aws disabled auto-merge June 12, 2026 21:14
Collapse over-explained docstrings to concise summaries, drop Args/Returns blocks for self-evident params, and shorten multi-line rationale comments to one line. No code changes.
@github-actions

Copy link
Copy Markdown
Contributor

Assessment: Comment (approve — pure documentation commit, no behavioral change)

Reviewed 97b70025 ("docs(memory): tighten verbose docstrings and comments") — 7 files, +208/−606.

I verified this is genuinely docs-only: I stripped all docstrings via AST from both 20a5a5e6 and 97b70025 and diffed the resulting executable code. All 7 files are byte-for-byte identical with docstrings removed (coordinator.py, model_extractor.py, triggers.py, extraction/types.py, memory_manager.py, types.py, exceptions.py). No logic, signatures, or control flow changed.

The trims condense some genuinely verbose prose (e.g. the multi-paragraph "how it works in three pieces" module docstring) into tighter summaries, while preserving the key concepts a reader needs — the per-store high-water mark, the per-store task chain serialization, backoff/probe behavior, and the never-raises guarantee. This is a net readability win.

Re-verified on 97b70025: 108 memory tests pass, ruff and mypy clean.

No code concerns. Position unchanged — the code is in good shape; the remaining open items are the design conversations in @pgrayy's review thread.

Remove class-enumerating module docstrings, self-evident Args/Returns
blocks on coordinator helpers, and an implementation note, keeping only
behavior callers cannot infer from signatures.
flush() also waits out saves that begin while it is waiting, a guarantee
callers cannot infer from the signature.
@github-actions

Copy link
Copy Markdown
Contributor

Assessment: Comment (approve — one-line docstring restore, no behavioral change)

Reviewed 16cf1e03 ("docs(memory): restore flush() concurrent-save note") — a single docstring line in coordinator.py.

This restores the "also waits out saves that start while waiting" note on flush(), which the previous trim had condensed away. Good follow-up — it brings back the one genuinely useful detail (that flush() re-checks for saves kicked off mid-wait) without reintroducing the verbose prose.

AST-strip confirms 9d8851c5 vs 16cf1e03 is identical executable code. Re-verified: 108 memory tests pass, ruff and mypy clean.

No concerns. Position unchanged.

@JackYPCOnline JackYPCOnline merged commit 713a6de into strands-agents:main Jun 15, 2026
43 of 56 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area-async Related to asynchronous flows or multi-threading area-hooks Features or requests that might be implementable via hooks area-persistence Session management or checkpointing enhancement New feature or request python Pull requests that update python code size/xl

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants